using SharpGraphT.Sharp.System.Collections.Generic;
using SharpGraphT.Sharp.Util;

namespace SharpGraphT.Graph.Specifics;

public class FastLookupDirectedSpecifics<TV, TE> : DirectedSpecifics<TV, TE>
    where TE : class, new() {
    private static readonly long SerialVersionUid = 4089085208843722263L;

    protected IDictionary<(TV, TV), ISet<TE>> touchingVerticesToEdgeMap;

    public FastLookupDirectedSpecifics(IGraph<TV, TE> graph,
        IDictionary<TV, DirectedEdgeContainer<TV, TE>> vertexMap,
        IDictionary<(TV, TV), ISet<TE>> touchingVerticesToEdgeMap,
        IEdgeSetFactory<TV, TE> edgeSetFactory) : base(graph, vertexMap,
        edgeSetFactory) =>
        this.touchingVerticesToEdgeMap = touchingVerticesToEdgeMap;

    public override IReadOnlySet<TE> GetAllEdges(TV sourceVertex,
        TV targetVertex) {
        if (!graph.ContainsVertex(sourceVertex) ||
            !graph.ContainsVertex(targetVertex)) {
            return null;
        }

        if (!touchingVerticesToEdgeMap.TryGetValue((sourceVertex, targetVertex),
                out var edges)) {
            return Collections.GetEmptyReadOnlySet<TE>();
        }

        var edgeSet = edgeSetFactory.CreateEdgeSet(sourceVertex);
        edgeSet.AddRange(edges);
        return edgeSet;
    }

    public override TE GetEdge(TV sourceVertex, TV targetVertex) =>
        touchingVerticesToEdgeMap.TryGetValue((sourceVertex, targetVertex),
            out var edges)
            ? edges.First()
            : null;

    public override bool AddEdgeToTouchingVertices(TV sourceVertex,
        TV targetVertex, TE e) {
        if (!base.AddEdgeToTouchingVertices(sourceVertex, targetVertex, e)) {
            return false;
        }

        AddToIndex(sourceVertex, targetVertex, e);
        return true;
    }

    public override bool AddEdgeToTouchingVerticesIfAbsent(TV sourceVertex,
        TV targetVertex, TE e) =>
        GetEdge(sourceVertex, targetVertex) == null &&
        AddEdgeToTouchingVertices(sourceVertex, targetVertex, e);

    public override TE CreateEdgeToTouchingVerticesIfAbsent(TV sourceVertex,
        TV targetVertex, Func<TE> edgeSupplier) {
        var edge = GetEdge(sourceVertex, targetVertex);
        if (edge != null) {
            return null;
        }

        var e = edgeSupplier();
        AddEdgeToTouchingVertices(sourceVertex, targetVertex, e);
        return e;
    }

    public override void RemoveEdgeFromTouchingVertices(TV sourceVertex,
        TV targetVertex, TE e) {
        base.RemoveEdgeFromTouchingVertices(sourceVertex, targetVertex, e);
        RemoveFromIndex(sourceVertex, targetVertex, e);
    }

    protected void AddToIndex(TV sourceVertex, TV targetVertex, TE e) {
        var vertexPair = (sourceVertex, targetVertex);
        if (touchingVerticesToEdgeMap.TryGetValue(vertexPair,
                out var edgeSet)) {
            edgeSet.Add(e);
        } else {
            edgeSet = edgeSetFactory.CreateEdgeSet(sourceVertex);
            edgeSet.Add(e);
            touchingVerticesToEdgeMap[vertexPair] = edgeSet;
        }
    }

    protected void RemoveFromIndex(TV sourceVertex, TV targetVertex, TE e) {
        var vertexPair = (sourceVertex, targetVertex);
        if (!touchingVerticesToEdgeMap.TryGetValue(vertexPair,
                out var edgeSet)) {
            return;
        }

        edgeSet.Remove(e);
        if (!edgeSet.Any()) {
            touchingVerticesToEdgeMap.Remove(vertexPair);
        }
    }
}